import torch
import torch.nn as nn


class DoubleConv(nn.Module):
    """(卷积 => ReLU) * 2"""

    def __init__(self, in_channels, out_channels):
        super().__init__()
        self.conv = nn.Sequential(
            nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1, stride=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1, stride=1),
            nn.ReLU(inplace=True),
        )

    def forward(self, x):
        return self.conv(x)


class Encoder(nn.Module):
    """编码器块：包含双卷积层和下采样"""

    def __init__(self, in_channels, out_channels):
        super().__init__()
        self.conv = DoubleConv(in_channels, out_channels)
        self.pool = nn.MaxPool2d(kernel_size=2, stride=2)

    def forward(self, x):
        x = self.conv(x)
        saved = x  # 保存特征图用于跳跃连接
        x = self.pool(x)
        return x, saved


class Decoder(nn.Module):
    """解码器块：包含上采样、特征拼接和双卷积"""

    def __init__(self, in_channels, out_channels, encoder_channels):
        super().__init__()
        self.up = nn.ConvTranspose2d(in_channels, out_channels, kernel_size=2, stride=2)
        self.conv = DoubleConv(out_channels + encoder_channels, out_channels)

    def forward(self, x, saved):
        x = self.up(x)
        x = torch.cat([x, saved], dim=1)  # 沿通道维度拼接
        x = self.conv(x)
        return x


class UNet(nn.Module):
    def __init__(self):
        super().__init__()
        # 编码器
        self.enc1 = Encoder(1, 64)  # 下采样后输出: 64x54x54
        self.enc2 = Encoder(64, 128)  # 下采样后输出: 128x27x27

        # 瓶颈层（无下采样）
        self.bottleneck = DoubleConv(128, 256)  # 输出: 256x27x27

        # 解码器
        self.dec1 = Decoder(256, 128, encoder_channels=128)  # 输出: 128x54x54
        self.dec2 = Decoder(128, 64, encoder_channels=64)  # 输出: 64x108x108

        # 最终1x1卷积
        self.final = nn.Conv2d(64, 1, kernel_size=1)

    def forward(self, x):
        # 编码器
        x, saved1 = self.enc1(x)  # saved1: 64x108x108 (跳跃连接1)
        x, saved2 = self.enc2(x)  # saved2: 128x54x54 (跳跃连接2)

        # 瓶颈层
        x = self.bottleneck(x)  # 256x27x27

        # 解码器
        x = self.dec1(x, saved2)  # 128x54x54
        x = self.dec2(x, saved1)  # 64x108x108

        # 最终输出
        return self.final(x)  # 1x108x108


# 验证模型
if __name__ == "__main__":
    model = UNet()
    input_tensor = torch.randn(8, 1, 108, 108)  # 输入尺寸: (batch, channel, H, W)
    output = model(input_tensor)
    print(f"输入尺寸: {input_tensor.shape}")
    print(f"输出尺寸: {output.shape}")
